import numpy as np
import scipy.sparse as sp

import lzyutil as util
from config import args


def LPSI_process(adj: np.ndarray, alpha, labels: np.ndarray,
                 normalized=False) -> np.ndarray:
    S = np.asmatrix(adj if normalized else util.get_norm_laplacian(adj, enhance_self_loop=False))
    numnodes = S.shape[0]
    assert labels.shape[0] == numnodes
    return np.array((1 - alpha) * (np.identity(numnodes) - alpha * S).I @ labels).flatten()


def LPSI_batch_iter_process(adj: sp.coo_matrix, alpha, labels: sp.coo_matrix,
                            normalized=False, iter_step=5) -> sp.coo_matrix:
    S = adj if normalized else util.get_norm_laplacian_sparse(adj, enhance_self_loop=True)
    ret = labels.copy()
    assert len(ret.shape) == 2 and ret.shape[1] == S.shape[0]
    # ret = sp.coo_matrix(labels)  # make a copy
    for _ in range(iter_step):
        ret = alpha * ret @ S + (1 - alpha) * ret
    return ret


def LPSI_batch_cached_process(lpsi_mat: sp.coo_matrix, labels: sp.coo_matrix) -> sp.coo_matrix:
    assert lpsi_mat is not None
    assert len(labels.shape) == 2 and labels.shape[1] == lpsi_mat.shape[0]
    return labels @ lpsi_mat


def LPSI_iter_process(adj: sp.coo_matrix, alpha, labels: sp.coo_matrix,
                      normalized=False, iter_step=5) -> np.ndarray:
    S = adj if normalized else util.get_norm_laplacian_sparse(adj, enhance_self_loop=True)
    ret = sp.coo_matrix(labels)  # make a copy
    for _ in range(iter_step):
        ret = alpha * S * ret + (1 - alpha) * labels
    return ret.toarray().flatten()


def extended_LPSI_cached_process(s_lpsi: sp.coo_matrix, labels: np.ndarray) -> np.ndarray:
    result = np.zeros((args.num_features, labels.shape[0]))
    result[0, :] = labels
    no_negative = np.ones_like(labels)
    no_positive = np.ones_like(labels)
    no_positive[labels > 0] = 0
    no_negative[labels < 0] = 0
    sparse_labels = sp.coo_matrix(labels.reshape((len(labels), 1)))
    sparse_nn = sp.coo_matrix(no_negative.reshape(len(labels), 1))
    sparse_np = sp.coo_matrix(no_positive.reshape(len(labels), 1))

    result[1, :] = LPSI_cached_process(s_lpsi, sparse_labels)
    result[2, :] = LPSI_cached_process(s_lpsi, sparse_nn)
    result[3, :] = LPSI_cached_process(s_lpsi, sparse_np)
    return result


def extended_LPSI_iter_process(adj: sp.coo_matrix, alpha, labels: np.ndarray) -> np.ndarray:
    result = np.zeros((args.num_features, labels.shape[0]))
    result[0, :] = labels
    no_negative = np.ones_like(labels)
    no_positive = np.ones_like(labels)
    no_positive[labels > 0] = 0
    no_negative[labels < 0] = 0
    sparse_labels = sp.coo_matrix(labels.reshape((len(labels), 1)))
    sparse_nn = sp.coo_matrix(no_negative.reshape(len(labels), 1))
    sparse_np = sp.coo_matrix(no_positive.reshape(len(labels), 1))
    result[1, :] = LPSI_iter_process(adj, alpha, sparse_labels, normalized=True)
    result[2, :] = LPSI_iter_process(adj, alpha, sparse_nn, normalized=True)
    result[3, :] = LPSI_iter_process(adj, alpha, sparse_np, normalized=True)
    return result


def LPSI_cached_process(s_lpsi: sp.coo_matrix, labels: sp.coo_matrix):
    assert s_lpsi is not None
    return s_lpsi.dot(labels).toarray().flatten()


if __name__ == '__main__':
    print(args.learning_rate)
